import pymysql
import pandas as pd
import time
from utils.log import MyLog
from tqdm import tqdm


class SqlHandler:
    def __init__(self, config):
        self.config = config
        self.keys = ['id','title', 'category', 'ctime', 'ext0']
        self.log = MyLog(config.log_dir,__name__).getlog()

    def __enter__(self):
        self.connection = pymysql.connect(host=self.config.host, port=self.config.port,
                                          database=self.config.database, charset=self.config.charset,
                                          user=self.config.user, password=self.config.password)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.connection.close()

    # def searchhistory(self):
    #     # sql = "select id, title, category, ctime,ext0 from sina_news_hot_search where ext1='1' order by ctime "
    #     data = self.history_search()
    #     data = data.dropna(subset=['title'])
    #     history = data.drop_duplicates(subset='title', keep='first')
    #     return history

    def search_by_words(self, words):
        if not words:
            return []
        keywords = words[-1] if len(words) == 1 else "|".join(words)
        sql = "select id, title, category, ctime,ext0 from sina_news_hot_search where ext1='1' and title regexp " + '\'' + keywords + '\''
        data = self.large_search(sql)
        data = data.drop_duplicates(subset='title', keep='first')
        history_data = [{key: values for key, values in zip(self.keys, line)} for line in data.values]
        return history_data

    def groupid_title(self, group_id):
        if not group_id: return None
        sql = f"select title from sina_news_hot_search where ext0 = {group_id} order by ctime desc limit 1"
        result = self.single_search(sql)
        return result[-1][0] if result else None

    def search_by_group_id(self, group_id):
        if not group_id: return {}
        sql = f"select id, title, category, ctime,ext0 from sina_news_hot_search where ext0 = {group_id} order by ctime desc limit 1"
        result = self.single_search(sql)
        print(result)
        return {k: v for k, v in zip(self.keys, list(result[0]))}

    def search_current(self):
        sql = f"select ext0 from sina_news_hot_search where ext0 != '' and ext1='0'order by ctime desc"
        # sql = f"select ext0 from sina_news_hot_search where ext0 != '' and ext1='0'order by ctime desc"
        groupid_list = self.single_search(sql)
        return groupid_list

    def single_search(self, sql):
        try:
            with self.connection.cursor() as cursor:
                cursor.execute(sql)
                return cursor.fetchall()
        except Exception as e:
            self.log.exception(e)
        finally:
            # self.connect.close()
            pass
    def searchhistory(self):
        history = pd.DataFrame(columns=self.keys)
        # count_sql = f'SELECT count(*) FROM sina_news_hot_search WHERE ext0 ="1"'
        id_range = 'SELECT min(id),max(id) FROM sina_news_hot_search'
        with self.connection.cursor() as cursor:
            try:
                cursor.execute(id_range)
                cur_begin,cur_end = cursor.fetchall()[0]
                step = (int(cur_end) - int(cur_begin))//1000
                for batch in tqdm(range(cur_begin,cur_end,step),desc='历史热榜下载'):
                    sql = f"select id, title, category, ctime,ext0 from sina_news_hot_search where ext1='1' and id >{batch} order by ctime limit {step}"
                    data = pd.read_sql(sql,self.connection)
                    data = data.dropna(subset=['title'])
                    history = pd.concat([history,data],axis=0)
            except Exception as e:
                self.log.exception(e)
        self.log.info(f'热榜去重前数据{history.shape}')
        history = history.dropna(subset=['title'])
        history = history.drop_duplicates(subset='title', keep='first')
        self.log.info(f'热榜去重后数据{history.shape}')
        return history


    def large_search(self, sql):
        try:
            return pd.read_sql(sql, self.connection)
        except Exception as e:
            self.log.exception(e)
        finally:
            # self.connect.close()
            pass


if __name__ == '__main__':
    from config import Config
    opt = Config()
    with SqlHandler(opt) as sq:
        res = sq.search_by_group_id('540644900176498688')
        print(res)

